/* Copyright  2009
   @Author
   Richard Changde Yin            e-mail yinchangde@hotmail.com

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; version 2 of the License.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA */
/*
 * bucket_tree.cc
 *
 *  Created on: 2009-6-26
 *  Update    : 2009-8-12
 */
#include<server_includes.h>
#include<index.h>
#include<allocator.h>
#include<database.h>
#include<databaseManager.h>
#include<databaseManagerImpl.h>
#include<debug.h>

#include <util.h>
#include <error.h>
#include <current_session.h>
using namespace BackEnd;

void BucketTree::init_tree(uint32_t default_alloc_size,
                           int size,
                           qsort_cmp2 i_compare,
                           void * _custom_arg)
{
  if (default_alloc_size < DEFAULT_ALLOC_SIZE)
    default_alloc_size= DEFAULT_ALLOC_SIZE;
  default_alloc_size= COLLAPSAR_ALIGN(default_alloc_size, DEFAULT_ALIGN_SIZE);
  memset(&null_element,0,sizeof(null_element));
  root= &null_element;
  compare=i_compare;
  size_of_element=size > 0 ? (uint32_t) size : 0;

  allocated=0;
  elements_in_tree=0;
  custom_arg = _custom_arg;
  null_element.colour=BLACK;
  null_element.left=null_element.right=0;
  flag= 0;
  offset_to_key=sizeof(BUCKET_TREE_ELEMENT);
  hashNodechunk  = 0 ;
  size_of_element+=sizeof(void*);
  mb_isInit = true ;
  return;
}

void BucketTree::free_tree( myf )
{
  if (root)    /* If initialized */
  {
    delete_tree_element(root);
  }
  root= &null_element;
  elements_in_tree=0;
  allocated=0;
  BKT_TREE_ARG* bk_tree_arg  = (BKT_TREE_ARG*)custom_arg ;
  delete bk_tree_arg;
  custom_arg = 0;
  hashNodechunk=0;
  mb_isInit = false ;
  return;
}

void BucketTree::delete_tree()
{
  this->free_tree(MYF(0));
}

void BucketTree::delete_tree_element(BUCKET_TREE_ELEMENT *element)
{
  assert(this->custom_arg) ; // put assert here.
  if (element != &null_element)
  {
    BKT_TREE_ARG* bk_tree_arg  = (BKT_TREE_ARG*)this->custom_arg ;
    delete_tree_element(element->left);

    BUCKET_TREE_ELEMENT* right_elem = element->right;

    Chunk* chunkptr =bk_tree_arg->chunkptr ;
    Database* dbptr =bk_tree_arg->dbptr    ;

    if(element->key_chunk && allocate_bychunk == bk_tree_arg->alloc_type) {
        DatabaseManager* dbmgr = current_session->getDatabaseManager() ;
        DatabaseManagerImpl* dbmgrimpl = (DatabaseManagerImpl*) dbmgr ;
        element->key_chunk->free(dbptr,element->key_ptr);
        dbmgrimpl->deleteUserIndexChunk((Chunk*)element->key_chunk);
    }
    // handle the sublist.
    if(hashNodechunk){
      if (element->bucketList_)
       element->bucketList_->removeAll((Chunk*)hashNodechunk,dbptr);
    }
    if (element->bucketList_) delete element->bucketList_;
    chunkptr->free(dbptr,element);
    delete_tree_element(right_elem);
  }
}


BUCKET_TREE_ELEMENT* BucketTree::tree_insert(void *key,
                                      uint32_t key_size,
                                      void* i_custom_arg,
                                      int&  rt)
{
  int cmp;
  BUCKET_TREE_ELEMENT *element,***parent;
  BKT_TREE_ARG* bk_tree_arg  = (BKT_TREE_ARG*)i_custom_arg ;

  parent  = parents;
  *parent = &root; element = root;
   for (;;)
  {
    if (element == &null_element ||
    (cmp = (*compare)(i_custom_arg, get_element_key(element),key)) == 0)
      break;
    if (cmp < 0)
    {
      *++parent= &element->right; element= element->right;
    }
    else
    {
      *++parent = &element->left; element= element->left;
    }
  } // End of for (;;)
  if (element == &null_element)
  {

    DbRetVal rv = OK;
    Chunk* chunkptr =bk_tree_arg->chunkptr ;
    Database* dbptr =bk_tree_arg->dbptr    ;
    element = (BUCKET_TREE_ELEMENT*) chunkptr->allocate(dbptr,&rv);
    if (!element)
    {
      rt= HA_ERR_OUT_OF_MEM;
      return (0);
    }

    **parent=element;
    element->left=element->right= &null_element;

    element->hashbucketList_=0;
    element->bucketList_    =0;
    BKT_TREE_ARG* basic_tree_arg  = (BKT_TREE_ARG*)this->custom_arg ;
    /* copy record ptr */
    if(basic_tree_arg->alloc_type == allocate_bychunk)
    {
      DatabaseManager* dbmgr = current_session->getDatabaseManager() ;
      DatabaseManagerImpl* dbmgrimpl = (DatabaseManagerImpl*) dbmgr ;
      /// Todo: get system database lock.
      Chunk* keychunk = dbmgrimpl->createUserIndexChunk(key_size);
      if(!keychunk)
      {
       rt= HA_ERR_OUT_OF_MEM;
       chunkptr->free(dbptr,element);
       return (0);
      }
      element->key_chunk = keychunk;
      element->key_ptr = keychunk->allocate(dbptr,&rv);
      if(!element->key_ptr){
        rt= HA_ERR_OUT_OF_MEM;
        dbmgrimpl->deleteUserIndexChunk((Chunk*)element->key_chunk);
        chunkptr->free(dbptr,element); return (0);
      }
      element->bucketList_ = new BucketList;
      uint32_t alloc_size=chunkptr->getSize()+keychunk->getSize();
      allocated+=alloc_size;
    }

    if(basic_tree_arg->alloc_type == assignment) element->key_ptr = key ;
    else
      memcpy((unsigned char*) element->key_ptr, key, key_size);
    element->count=1;    /* May give warning in purify */
    elements_in_tree++;
    this->rb_insert(parent,element);/* re-balance tree */
  }
  else
  {
    if (flag & TREE_NO_DUPS){
      rt = HA_ERR_FOUND_DUPP_KEY;
      return(0);
    }
    element->count++;
    /* Avoid a wrap over of the count. */
    if (! element->count)
      element->count--;
  }
  return element;
}



int BucketTree::tree_delete(void *key, uint32_t ,
                            void *i_custom_arg)
{
  int cmp,remove_colour;
  BUCKET_TREE_ELEMENT *element,***parent, ***org_parent, *nod;
  BKT_TREE_ARG* bk_tree_arg     = (BKT_TREE_ARG*)i_custom_arg ;
  BKT_TREE_ARG* basic_tree_arg  = (BKT_TREE_ARG*)this->custom_arg ;
  parent= parents;
  *parent= &root; element= root;
  for (;;)
  {
    if (element == &null_element)
      return 1;                /* Was not in tree */

    if ((cmp = (*compare)(i_custom_arg, get_element_key(element),
                                        key)) == 0)
      break;
    if (cmp < 0)
    {
      *++parent= &element->right; element= element->right;
    }
    else
    {
      *++parent = &element->left; element= element->left;
    }
  }
  if (element->left == &null_element)
  {
    (**parent)=element->right;
    remove_colour= element->colour;
  }
  else if (element->right == &null_element)
  {
    (**parent)=element->left;
    remove_colour= element->colour;
  }
  else
  {
    org_parent= parent;
    *++parent= &element->right; nod= element->right;
    while (nod->left != &null_element)
    {
      *++parent= &nod->left; nod= nod->left;
    }
    (**parent)=nod->right;      /* unlink nod from tree */
    remove_colour= nod->colour;
    org_parent[0][0]=nod;       /* put y in place of element */
    org_parent[1]= &nod->right;
    nod->left=element->left;
    nod->right=element->right;
    nod->colour=element->colour;
  }
  if (remove_colour == BLACK)
      rb_delete_fixup(parent);
  Chunk* chunkptr =bk_tree_arg->chunkptr ;
  Database* dbptr =bk_tree_arg->dbptr    ;
  allocated-= chunkptr->getSize();
  // handle the sublist.
  if(hashNodechunk ){
    if (element->bucketList_)
        element->bucketList_->removeAll((Chunk*)hashNodechunk,dbptr);
  }
  if(element->key_chunk && basic_tree_arg->alloc_type == allocate_bychunk ) {
      DatabaseManager* dbmgr = current_session->getDatabaseManager() ;
      DatabaseManagerImpl* dbmgrimpl = (DatabaseManagerImpl*) dbmgr ;
      allocated-= element->key_chunk->getSize();
      element->key_chunk->free(dbptr,element->key_ptr);
      dbmgrimpl->deleteUserIndexChunk((Chunk*)element->key_chunk);
  }
  if (element->bucketList_)delete element->bucketList_;
  chunkptr->free(dbptr,element);
  elements_in_tree--;
  return 0;
}


void* BucketTree::tree_search(void *key, void *i_custom_arg)
{
  int cmp;
  BKT_TREE_ARG* _arg = (BKT_TREE_ARG*)i_custom_arg;
  BUCKET_TREE_ELEMENT *element=root;

  for (;;)
  {
    if (element == &null_element)
      return (void*) 0;
    if ((cmp = (*compare)(i_custom_arg, get_element_key(element),
                                key)) == 0)
    {
     if(_arg->search_type == get_key)
       return get_element_key(element);
     else return element ;
    }
    if (cmp < 0)
      element=element->right;
    else
      element=element->left;
  }
}


void* BucketTree::tree_search_key(
                      const void *key,
                      BUCKET_TREE_ELEMENT **i_parents,
                      BUCKET_TREE_ELEMENT ***last_pos,
                      enum ha_rkey_function i_flag,
                      void *i_custom_arg)
{
  int cmp;
  BUCKET_TREE_ELEMENT *element= root;
  BUCKET_TREE_ELEMENT **last_left_step_parent= NULL, **last_right_step_parent= NULL;
  BUCKET_TREE_ELEMENT **last_equal_element= NULL;
  BKT_TREE_ARG* _arg = (BKT_TREE_ARG*)i_custom_arg;
/*
  TODO: support for HA_READ_KEY_OR_PREV, HA_READ_PREFIX flags if needed.
*/
  *i_parents = &null_element;
  while (element != &null_element)
  {
    *++i_parents= element;
    if ((cmp= (*compare)(_arg, get_element_key(element),key)) == 0)

    {
      switch (i_flag) {
      case HA_READ_KEY_EXACT:
      case HA_READ_KEY_OR_NEXT:
      case HA_READ_BEFORE_KEY:
        last_equal_element= i_parents;
        cmp= 1;
        break;
      case HA_READ_AFTER_KEY:
        cmp= -1;
      break;
      case HA_READ_PREFIX_LAST:
      case HA_READ_PREFIX_LAST_OR_PREV:
         last_equal_element= i_parents;
         cmp= -1;
      break;
      default:
       return NULL;
      }
    }
    if (cmp < 0) /* element < key */
    {
      last_right_step_parent= i_parents;
      element= element->right;
    }
    else
    {
      last_left_step_parent= i_parents;
      element= element->left;
    }
  }
  switch (flag) {
  case HA_READ_KEY_EXACT:
  case HA_READ_PREFIX_LAST:
    *last_pos= last_equal_element;
    break;
  case HA_READ_KEY_OR_NEXT:
    *last_pos= last_equal_element ? last_equal_element : last_left_step_parent;
    break;
  case HA_READ_AFTER_KEY:
    *last_pos= last_left_step_parent;
    break;
  case HA_READ_PREFIX_LAST_OR_PREV:
    *last_pos= last_equal_element ? last_equal_element : last_right_step_parent;
    break;
  case HA_READ_BEFORE_KEY:
    *last_pos= last_right_step_parent;
    break;
  default:
    return 0;
  }
  if (_arg->search_type == get_element)
    { return *last_pos ? (**last_pos) : NULL; }
  else
    { return *last_pos ? get_element_key(**last_pos)  : NULL;}
}

/*
  Search first (the most left) or last (the most right) tree element
*/
void *BucketTree::tree_search_edge(BUCKET_TREE_ELEMENT **i_parents,
                                   BUCKET_TREE_ELEMENT ***last_pos,
                                   int child_offs)
{
  BUCKET_TREE_ELEMENT *element= root;

  *i_parents= &null_element;
  while (element != &null_element)
  {
    *++i_parents= element;
    element= BUCKET_ELEMENT_CHILD(element, child_offs);
  }
  *last_pos= i_parents;
  return **last_pos != &null_element?get_element_key(**last_pos) : NULL;

}

void *BucketTree::tree_search_next(BUCKET_TREE_ELEMENT ***last_pos,
                                  int l_offs,
                                  int r_offs)
{
  BUCKET_TREE_ELEMENT *x= **last_pos;

  if (BUCKET_ELEMENT_CHILD(x, r_offs) != &null_element)
  {
    x= BUCKET_ELEMENT_CHILD(x, r_offs);
    *++*last_pos= x;
    while (BUCKET_ELEMENT_CHILD(x, l_offs) != &null_element)
    {
      x= BUCKET_ELEMENT_CHILD(x, l_offs);
      *++*last_pos= x;
    }
    return get_element_key(x);
  }
  else
  {
    BUCKET_TREE_ELEMENT *y= *--*last_pos;
    while (y != &null_element && x == BUCKET_ELEMENT_CHILD(y, r_offs))
    {
      x= y;
      y= *--*last_pos;
    }
    return y == &null_element ? NULL : get_element_key (y);
  }
}

/*
  Expected that tree is fully balanced
  (each path from root to leaf has the same length)
*/
ha_rows BucketTree::tree_record_pos(const void *key,
                        enum ha_rkey_function i_flag, void *i_custom_arg)
{
  int cmp;
  BUCKET_TREE_ELEMENT *element= root;
  double left= 1;
  double right= elements_in_tree;

  while (element != &null_element)
  {
    if ((cmp= (*compare)(i_custom_arg, get_element_key( element),
                               key)) == 0)
    {
      switch (i_flag) {
      case HA_READ_KEY_EXACT:
      case HA_READ_BEFORE_KEY:
        cmp= 1;
        break;
      case HA_READ_AFTER_KEY:
        cmp= -1;
        break;
      default:
        return HA_POS_ERROR;
      }
    }
    if (cmp < 0) /* element < key */
    {
      element= element->right;
      left= (left + right) / 2;
    }
    else
    {
      element= element->left;
      right= (left + right) / 2;
    }
  }
  switch (flag) {
  case HA_READ_KEY_EXACT:
  case HA_READ_BEFORE_KEY:
    return (ha_rows) right;
  case HA_READ_AFTER_KEY:
    return (ha_rows) left;
  default:
    return HA_POS_ERROR;
  }
}



int BucketTree::tree_walk( tree_walk_action action,
                           void *argument, TREE_WALK visit)
{
  switch (visit) {
  case left_root_right:
    return tree_walk_left_root_right(root,action,argument);
  case right_root_left:
    return tree_walk_right_root_left(root,action,argument);
  }
  return 0; /* Keep gcc happy */
}



int BucketTree::tree_walk_left_root_right(BUCKET_TREE_ELEMENT *element,
                                          tree_walk_action action,
                                          void *argument)
{
  int error;
  if (element->left)  /* Not null_element */
  {
    if ((error=tree_walk_left_root_right(element->left,action,
                                        argument)) == 0 &&
              (error=(*action)(get_element_key(element),
               (element_count) element->count,
               argument)) == 0)
      error=tree_walk_left_root_right(element->right,action,argument);
    return error;
  }
  return 0;
}

int BucketTree::tree_walk_right_root_left(BUCKET_TREE_ELEMENT *element,
                              tree_walk_action action, void *argument)
{
  int error;
  if (element->right)/* Not null_element */
  {
    if ((error=tree_walk_right_root_left(element->right,action, argument)) == 0
              &&(error=(*action)(get_element_key(element),
                      (element_count) element->count,argument)) == 0)
     error=tree_walk_right_root_left(element->left,action,argument);
    return error;
  }
  return 0;
}


 /* Functions to fix up the tree after insert and delete */

void BucketTree::left_rotate(BUCKET_TREE_ELEMENT **parent, BUCKET_TREE_ELEMENT *leaf)
{
  BUCKET_TREE_ELEMENT *y;

  y=leaf->right;
  leaf->right=y->left;
  parent[0]=y;
  y->left=leaf;
}

void BucketTree::right_rotate(BUCKET_TREE_ELEMENT **parent, BUCKET_TREE_ELEMENT *leaf)
{
  BUCKET_TREE_ELEMENT *x;

  x=leaf->left;
  leaf->left=x->right;
  parent[0]=x;
  x->right=leaf;
}

void BucketTree::rb_insert(BUCKET_TREE_ELEMENT ***parent, BUCKET_TREE_ELEMENT *leaf)
{
  BUCKET_TREE_ELEMENT *y,*par,*par2;

  leaf->colour=RED;
  while (leaf != root && (par=parent[-1][0])->colour == RED)
  {
    if (par == (par2=parent[-2][0])->left)
    {
      y= par2->right;
      if (y->colour == RED)
      {
        par->colour=BLACK;
        y->colour=BLACK;
        leaf=par2;
        parent-=2;
        leaf->colour=RED; /* And the loop continues */
      }
      else
      {
         if (leaf == par->right)
        {
          left_rotate(parent[-1],par);
          par=leaf;       /* leaf is now parent to old leaf */
        }
        par->colour=BLACK;
        par2->colour=RED;
        right_rotate(parent[-2],par2);
        break;
      }
    }
    else
    {
      y= par2->left;
      if (y->colour == RED)
      {
        par->colour=BLACK;
        y->colour=BLACK;
        leaf=par2;
        parent-=2;
        leaf->colour=RED; /* And the loop continues */
      }
      else
      {
         if (leaf == par->left)
        {
          right_rotate(parent[-1],par);
          par=leaf;
        }
        par->colour=BLACK;
        par2->colour=RED;
        left_rotate(parent[-2],par2);
        break;
      }
    }
  }
  root->colour=BLACK;
}

void BucketTree::rb_delete_fixup(BUCKET_TREE_ELEMENT ***parent)
{
  BUCKET_TREE_ELEMENT *x,*w,*par;

  x= **parent;
  while (x != root && x->colour == BLACK)
  {
    if (x == (par=parent[-1][0])->left)
    {
      w=par->right;
      if (w->colour == RED)
      {
        w->colour=BLACK;
        par->colour=RED;
        left_rotate(parent[-1],par);
        parent[0]= &w->left;
        *++parent= &par->left;
        w=par->right;
      }
      if (w->left->colour == BLACK && w->right->colour == BLACK)
      {
        w->colour=RED;
        x=par;
        parent--;
      }
      else
      {
        if (w->right->colour == BLACK)
       {
         w->left->colour=BLACK;
         w->colour=RED;
         right_rotate(&par->right,w);
         w=par->right;
        }
       w->colour=par->colour;
       par->colour=BLACK;
       w->right->colour=BLACK;
       left_rotate(parent[-1],par);
       x=root;
       break;
      }
    }
    else
    {
      w=par->left;
      if (w->colour == RED)
      {
        w->colour=BLACK;
        par->colour=RED;
        right_rotate(parent[-1],par);
        parent[0]= &w->right;
        *++parent= &par->right;
        w=par->left;
      }
      if (w->right->colour == BLACK && w->left->colour == BLACK)
      {
        w->colour=RED;
        x=par;
        parent--;
      }
      else
      {
        if (w->left->colour == BLACK)
       {
          w->right->colour=BLACK;
          w->colour=RED;
          left_rotate(&par->left,w);
          w=par->left;
       }
       w->colour=par->colour;
       par->colour=BLACK;
       w->left->colour=BLACK;
       right_rotate(parent[-1],par);
       x=root;
       break;
      }
    }
  }
  x->colour=BLACK;
}

